import os

from ultralytics import YOLO


def train():
    # 获得当前文件的路径作为基路径
    curPath = os.getcwd()
    print(curPath)

    # 从YAML建立并转移权重
    model = YOLO("./yolov8_trainOriginal.yaml").load("./yolov8n.pt")

    # 训练模型,这里会自动在当前目录下找yolov8n.pt,没有则自动下载
    model.train(
        data="./ZTestYOLOv8.yaml",
        # data="./ZTestYOLOv8_special.yaml",
        epochs=200,
        # batch=32,
        # workers=28,
        # device=[0,1],
        # cache=False,
    )


if __name__ == "__main__":
    train()
